/**create by liuhua at 2019年5月22日 上午9:14:38**/
package com.star.truffle.gateway.filter;

import java.net.URI;
import java.nio.CharBuffer;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.commons.lang.StringUtils;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;

import io.netty.buffer.ByteBufAllocator;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

@Slf4j
public class PublicGlobalFilter implements GlobalFilter, Ordered {

  private static final String REQUEST_TIME_BEGIN = "requestTimeBegin";
  private final ObjectMapper mapper = new ObjectMapper();

  @Override
  public int getOrder() {
    return -999;
  }

  @Override
  public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
    exchange.getAttributes().put(REQUEST_TIME_BEGIN, System.currentTimeMillis());

    String ip = "";
    List<String> ips = exchange.getRequest().getHeaders().get("X-Real-IP");
    if(null != ips && ips.size() > 0) {
      ip = ips.get(0);
    }
    if(StringUtils.isBlank(ip)) {
      ip = "127.0.0.1";
    }
    String token = exchange.getRequest().getHeaders().getFirst("token");
    log.info("ip->{}, token->{}", ip, token);
//    if (StringUtils.isNotBlank(token)) {
      Long userId = null;
      return modifyBody(exchange, chain, userId, ip);
//    } else {
//      return chain.filter(exchange).then(runnable(exchange, exchange.getRequest()));
//    }
  }

  private Mono<Void> modifyBody(ServerWebExchange exchange, GatewayFilterChain chain, Long userId, String ip) {
    ServerHttpRequest serverHttpRequest = exchange.getRequest();
    HttpMethod method = serverHttpRequest.getMethod();
    String contentType = serverHttpRequest.getHeaders().getFirst(HttpHeaders.CONTENT_TYPE);
    URI uri = serverHttpRequest.getURI();
    
    if (method == HttpMethod.POST && (MediaType.APPLICATION_FORM_URLENCODED_VALUE.equalsIgnoreCase(contentType) || MediaType.APPLICATION_JSON_VALUE.equalsIgnoreCase(contentType))) {
      Object requestBody = exchange.getAttribute("cachedRequestBodyObject");
      String bodyStr = String.valueOf(requestBody);
      log.info("bodyStr before->{}", bodyStr);
      if (MediaType.APPLICATION_FORM_URLENCODED_VALUE.equalsIgnoreCase(contentType)) {
        if(null != userId) {
          bodyStr = String.format(bodyStr + "&userId=%s", userId);
        }
        if(StringUtils.isNotBlank(ip)) {
          bodyStr = String.format(bodyStr + "&ip=%s", ip);
        }
      }
      if (MediaType.APPLICATION_JSON_VALUE.equalsIgnoreCase(contentType)) {
        try {
          if(StringUtils.isBlank(bodyStr)) {
            bodyStr = "{}";
          }
          Map<String, Object> params = mapper.readValue(bodyStr, new TypeReference<Map<String, Object>>() {});
          params.put("userId", String.valueOf(userId));
          params.put("ip", ip);
          bodyStr = mapper.writeValueAsString(params);
        } catch (Exception e) {
          log.error(e.getMessage(), e);
        }
      }
      log.info("bodyStr after->{}", bodyStr);
      DataBuffer bodyDataBuffer = stringBuffer(bodyStr);
      Flux<DataBuffer> bodyFlux = Flux.just(bodyDataBuffer);
      URI newUri = UriComponentsBuilder.fromUri(uri).build(true).toUri();
      ServerHttpRequest request = exchange.getRequest().mutate().uri(newUri).build();

      HttpHeaders headers = new HttpHeaders();
      headers.putAll(exchange.getRequest().getHeaders());
      headers.set("star", "666666");

      int length = bodyStr.getBytes().length;
      headers.remove(HttpHeaders.CONTENT_LENGTH);
      headers.setContentLength(length);

      if (StringUtils.isNotBlank(contentType)) {
        headers.set(HttpHeaders.CONTENT_TYPE, contentType);
      }
      request = new ServerHttpRequestDecorator(request) {
        @Override
        public HttpHeaders getHeaders() {
          long contentLength = headers.getContentLength();
          HttpHeaders httpHeaders = new HttpHeaders();
          httpHeaders.putAll(headers);
          if (contentLength > 0) {
            httpHeaders.setContentLength(contentLength);
          } else {
            httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
          }
          return httpHeaders;
        }

        @Override
        public Flux<DataBuffer> getBody() {
          return bodyFlux;
        }
      };
      request.mutate().header(HttpHeaders.CONTENT_LENGTH, Integer.toString(bodyStr.length()));
      return chain.filter(exchange.mutate().request(request).build()).then(runnable(exchange, exchange.getRequest()));
    } else if (method == HttpMethod.GET) {
      StringBuilder query = new StringBuilder();
      String originalQuery = uri.getRawQuery();
      if (StringUtils.isNotBlank(originalQuery)) {
        query.append(originalQuery);
        if (! originalQuery.endsWith("&")) {
          query.append("&");
        }
      }else {
        query.append("?");
      }
      query.append("userId=" + userId + "&ip=" + ip);
      URI newUri = UriComponentsBuilder.fromUri(uri).replaceQuery(query.toString()).build(true).toUri();
      ServerHttpRequest request = exchange.getRequest().mutate().uri(newUri).build();
      return chain.filter(exchange.mutate().request(request).build()).then(runnable(exchange, exchange.getRequest()));
    }
    return chain.filter(exchange).then(runnable(exchange, exchange.getRequest()));
  }

  private DataBuffer stringBuffer(String value) {
    byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
    NettyDataBufferFactory nettyDataBufferFactory = new NettyDataBufferFactory(ByteBufAllocator.DEFAULT);
    DataBuffer buffer = nettyDataBufferFactory.allocateBuffer(bytes.length);
    buffer.write(bytes);
    return buffer;
  }

  @SuppressWarnings("unused")
  private String resolveBodyFromRequest(ServerHttpRequest serverHttpRequest) {
    // 获取请求体
    Flux<DataBuffer> body = serverHttpRequest.getBody();

    AtomicReference<String> bodyRef = new AtomicReference<>();
    body.subscribe(buffer -> {
      CharBuffer charBuffer = StandardCharsets.UTF_8.decode(buffer.asByteBuffer());
      DataBufferUtils.release(buffer);
      bodyRef.set(charBuffer.toString());
    });
    // 获取request body
    return bodyRef.get();
  }

  private Mono<Void> runnable(ServerWebExchange exchange, ServerHttpRequest request, String... ext) {
    return Mono.fromRunnable(() -> {
      Long startTime = exchange.getAttribute(REQUEST_TIME_BEGIN);
      if (startTime != null) {
        StringBuilder sb = new StringBuilder(request.getURI().getRawPath()).append(": ").append(System.currentTimeMillis() - startTime).append("ms");
        sb.append(" params:").append(request.getQueryParams());
        if (null != ext && ext.length > 0) {
          sb.append(" ExtensionInfo:");
          for (String str : ext) {
            sb.append(" " + str);
          }
        }
        log.info(sb.toString());
      }
    });
  }
}
